from pro_authformer.src.main.util_tools.tool import load_config
from pro_authformer.src.main.servlet_main.f_Dataset import data_processor, authformer_dataset
from pro_authformer.src.main.util_tools.logging_utils import setup_logger
from pro_authformer.src.main.servlet_main.config import config_dict
from pro_authformer.src.main.servlet_main.f_trainer import Trainer
from pro_authformer.src.main.servlet_main.f_ModelManager import ModelManager


class AU:
    """
    功能：协调训练（Training）、验证（Verification）、评估（Evaluation）流程，串联数据处理、模型管理、工具模块。
    """
    def __init__(self) -> None:
        # 初始化日志记录器
        self.logger = setup_logger()
        # 初始化数据处理器
        self.data_processor = data_processor()
        # 初始化模型管理器
        self.model_manager = ModelManager()
        # 初始化训练器
        self.trainer = Trainer()
        
        
    def run(self, mode):
        # 流程入口，驱动训练、验证、评估的执行逻辑（可能支持参数配置或模式选择）
        if mode == "train":
            self.training()  # 执行训练流程
            self.verify()    # 训练后验证（可选）
        elif mode == "eval":
            self.eval()      # 独立评估（如测试集评估）
        elif mode == "verify":
            self.verify()    # 验证集验证（如训练中实时验证）
        # 收尾操作（如关闭日志、保存最终状态）
        self.cleanup()
        
    def training(self):
        # 1. 数据准备
        ann_info = self.data_processor.gen_ann_info(config_dict["data_path"])
        data = self.data_processor.split_data(ann_info)
        
        train_loader = authformer_dataset(data, mode, shuffle=True)
        val_loader = authformer_dataset(data, mode, shuffle=False)
        
        # 2. 模型初始化
        model = self.model_manager.load_model()
        
        # 3. 训练配置
        loss, optimizer, scheduler = self.trainer.train_setup()
        
        # 4. 迭代训练
        self.trainer.fit_one_epoch(model, train_loader, val_loader, loss, optimizer, scheduler)
        
        
    def verify(self):
        pass
    def eval(self):
        pass
    def cleanup(self):
        # 清理资源、保存模型等
        pass
    
if __name__ == "__main__":
    # 测试代码
    
    mode = "train"  # 模式选择
    config = load_config(mode)
    
    au = AU()
    au.run(mode)
